#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import re
import json
from datetime import datetime
from InfluxDB_setting.api.parse_lynis import parse_lynis_log
from langchain_ollama import OllamaLLM

def extract_metadata(log_content):
    """从日志内容中提取元数据"""
    metadata = {
        "os": "Linux",
        "lynis_version": "3.1.4",
        "audit_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "hostname": "localhost"
    }
    
    # 提取系统信息
    os_match = re.search(r'Operating system version:\s+(.+)', log_content)
    if os_match:
        metadata["os"] = os_match.group(1).strip()
    
    version_match = re.search(r'Program version:\s+(\d+\.\d+\.\d+)', log_content)
    if version_match:
        metadata["lynis_version"] = version_match.group(1)
    
    hostname_match = re.search(r'Hostname:\s+(.+)', log_content)
    if hostname_match:
        metadata["hostname"] = hostname_match.group(1).strip()
    
    return metadata

def calculate_security_score(vulnerabilities):
    """计算安全评分"""
    total_vulnerabilities = len(vulnerabilities)
    high_risk_count = sum(1 for v in vulnerabilities if "password" in v.get("loopAction", "").lower() or "ssh" in v.get("loopAction", "").lower())
    medium_risk_count = sum(1 for v in vulnerabilities if "update" in v.get("loopAction", "").lower() or "configure" in v.get("loopAction", "").lower())
    
    # 基础分数100分，每个漏洞扣分
    base_score = 100
    deduction = (high_risk_count * 5) + (medium_risk_count * 2) + (total_vulnerabilities - high_risk_count - medium_risk_count)
    final_score = max(0, base_score - deduction)
    
    # 评分描述
    if final_score >= 80:
        description = "优秀"
    elif final_score >= 60:
        description = "良好"
    elif final_score >= 40:
        description = "一般"
    else:
        description = "需要立即整改"
    
    return {
        "hardening_index": final_score,
        "description": description,
        "failed_tests": total_vulnerabilities,
        "high_risk": high_risk_count,
        "medium_risk": medium_risk_count
    }

def classify_vulnerabilities(vulnerabilities):
    """分类漏洞风险等级"""
    high_risk = []
    medium_risk = []
    low_risk = []
    
    for vuln in vulnerabilities:
        content = vuln.get("loopAction", "").lower()
        if any(keyword in content for keyword in ["password", "ssh", "root", "firewall", "certificate", "expire"]):
            high_risk.append(vuln)
        elif any(keyword in content for keyword in ["update", "disable", "configure", "harden"]):
            medium_risk.append(vuln)
        else:
            low_risk.append(vuln)
    
    return {
        "high_risk": high_risk,
        "medium_risk": medium_risk,
        "low_risk": low_risk
    }

def generate_llm_report(vulnerabilities, log_content):
    """使用LLM生成安全审计报告"""
    
    # 提取元数据
    metadata = extract_metadata(log_content)
    
    # 计算安全评分
    security_score = calculate_security_score(vulnerabilities)
    
    # 分类漏洞
    classified_vulns = classify_vulnerabilities(vulnerabilities)
    
    # 准备LLM输入数据
    llm_data = {
        "metadata": metadata,
        "security_score": security_score,
        "vulnerabilities": {
            "high_risk": classified_vulns["high_risk"],
            "medium_risk": classified_vulns["medium_risk"],
            "low_risk": classified_vulns["low_risk"]
        },
        "total_vulnerabilities": len(vulnerabilities)
    }
    
    # 构建LLM提示词
    prompt = f"""你作为银河麒麟系统安全专家，请基于以下JSON数据生成中文安全审计报告：

JSON数据：
{json.dumps(llm_data, ensure_ascii=False, indent=2)}

报告要求：
1. 语言：专业、简洁的中文技术报告
2. 结构：
一、审计概览
系统版本：{metadata['os']}
Lynis版本：{metadata['lynis_version']}
审计时间：{metadata['audit_date']}
安全评分：{security_score['hardening_index']}/100（{security_score['description']}）
失败检测项：{security_score['failed_tests']}项

二、关键漏洞分析
高危漏洞（风险≥8）
{{循环输出高危漏洞}}
中危漏洞（风险5-7）
{{循环输出中危漏洞}}

三、修复建议
紧急修复项
{{含证书过期和关键配置的修复建议}}
加固建议
{{其他安全加固建议}}

四、合规性状态
PCI-DSS合规：{{根据漏洞自动判断}}
等保要求：{{根据漏洞自动判断}}

五、行动路线图
title 安全修复计划
现在 ： 处理证书过期和高危漏洞
3天内 ： 修复中危漏洞
1周内 ： 完成合规性整改

请根据提供的漏洞数据生成详细的安全审计报告，重点关注高危漏洞的修复建议和合规性要求。

报告编写人：SmartCare_for_KylinOS"""

    try:
        # 初始化LLM
        llm = OllamaLLM(base_url="http://10.11.5.14:11434", model="deepseek-r1:14b")
        
        # 调用LLM生成报告
        response = llm.invoke(prompt)
        
        # 清理LLM输出中的思考内容
        cleaned_response = clean_llm_output(response)
        
        return cleaned_response
        
    except Exception as e:
        print(f"LLM调用失败: {e}")
        # 如果LLM调用失败，生成备用报告
        return generate_fallback_report(vulnerabilities, metadata, security_score, classified_vulns)

def generate_fallback_report(vulnerabilities, metadata, security_score, classified_vulns):
    """生成备用报告（当LLM调用失败时使用）"""
    print("LLM调用失败，生成备用报告...")
    report = f"""# 系统安全审计报告

**报告编写人：SmartCare_for_KylinOS**

## 一、审计概览

- **系统版本**: {metadata['os']}
- **Lynis版本**: {metadata['lynis_version']}
- **审计时间**: {metadata['audit_date']}
- **安全评分**: {security_score['hardening_index']}/100（{security_score['description']}）
- **失败检测项**: {security_score['failed_tests']}项
- **高危漏洞**: {len(classified_vulns['high_risk'])}项
- **中危漏洞**: {len(classified_vulns['medium_risk'])}项

## 二、关键漏洞分析

### 高危漏洞（风险≥8）

"""
    
    # 添加高危漏洞
    for vuln in classified_vulns['high_risk']:
        report += f"- **{vuln['loopName']}**: {vuln['loopAction']}\n"
    
    report += "\n### 中危漏洞（风险5-7）\n\n"
    
    # 添加中危漏洞
    for vuln in classified_vulns['medium_risk']:
        report += f"- **{vuln['loopName']}**: {vuln['loopAction']}\n"
    
    report += "\n## 三、修复建议\n\n### 紧急修复项\n\n"
    
    # 生成紧急修复建议
    urgent_items = []
    for vuln in classified_vulns['high_risk']:
        content = vuln.get("loopAction", "").lower()
        if "password" in content:
            urgent_items.append("配置密码策略，包括最小长度、复杂度要求和过期时间")
        elif "ssh" in content:
            urgent_items.append("加固SSH配置，禁用root登录，更改默认端口")
        elif "firewall" in content:
            urgent_items.append("配置防火墙规则，只允许必要的网络连接")
        elif "certificate" in content or "expire" in content:
            urgent_items.append("更新过期的证书和密钥")
    
    for item in set(urgent_items):  # 去重
        report += f"- {item}\n"
    
    report += "\n### 加固建议\n\n"
    
    # 生成加固建议
    hardening_items = []
    for vuln in classified_vulns['medium_risk']:
        content = vuln.get("loopAction", "").lower()
        if "update" in content:
            hardening_items.append("定期更新系统和软件包")
        elif "disable" in content:
            hardening_items.append("禁用不必要的服务和功能")
        elif "configure" in content:
            hardening_items.append("优化系统配置参数")
    
    for item in set(hardening_items):  # 去重
        report += f"- {item}\n"
    
    report += "\n## 四、合规性状态\n\n"
    
    # 分析合规性
    failed_items = []
    for vuln in vulnerabilities:
        content = vuln.get("loopAction", "").lower()
        if "password" in content:
            failed_items.append("密码策略")
        if "ssh" in content:
            failed_items.append("SSH安全配置")
        if "firewall" in content:
            failed_items.append("防火墙配置")
        if "logging" in content:
            failed_items.append("日志记录")
    
    pci_status = "部分合规" if len(failed_items) < 5 else "不合规"
    dengbao_status = "需要整改" if len(failed_items) > 0 else "基本合规"
    
    report += f"- **PCI-DSS合规**: {pci_status}（失败条款：{', '.join(set(failed_items)) if failed_items else '无'}）\n"
    report += f"- **等保要求**: {dengbao_status}（失败条款：{', '.join(set(failed_items)) if failed_items else '无'}）\n"
    
    report += "\n## 五、行动路线图\n\n"
    report += "### 安全修复计划\n\n"
    report += "- **现在**: 处理证书过期和高危漏洞\n"
    report += "- **3天内**: 修复中危漏洞\n"
    report += "- **1周内**: 完成合规性整改\n"
    
    return report

def clean_llm_output(llm_response):
    """清理LLM输出中的思考内容"""
    import re
    
    # 移除<think>标签及其内容
    cleaned_response = re.sub(r'<think>.*?</think>', '', llm_response, flags=re.DOTALL)
    
    # 移除思考内容（通常以"嗯，"、"首先，"等开头的大段文字）
    cleaned_response = re.sub(r'嗯，.*?(?=\n\n|\n#|\n##|\n###)', '', cleaned_response, flags=re.DOTALL)
    cleaned_response = re.sub(r'首先，.*?(?=\n\n|\n#|\n##|\n###)', '', cleaned_response, flags=re.DOTALL)
    
    # 移除AI思考的提示词
    cleaned_response = re.sub(r'我现在需要帮用户.*?(?=\n\n|\n#|\n##|\n###)', '', cleaned_response, flags=re.DOTALL)
    
    # 清理多余的空行
    cleaned_response = re.sub(r'\n\s*\n\s*\n', '\n\n', cleaned_response)
    
    # 确保报告以标题开始
    if not cleaned_response.strip().startswith('#'):
        # 查找第一个标题
        title_match = re.search(r'#\s*(.*?)(?:\n|$)', cleaned_response)
        if title_match:
            # 移除标题前的所有内容
            cleaned_response = re.sub(r'^.*?(?=#)', '', cleaned_response, flags=re.DOTALL)
    
    # 移除主标题、编写人信息和分隔线
    cleaned_response = re.sub(r'^#\s*银河麒麟系统安全审计报告\s*\n\*\*报告编写人：SmartCare_for_KylinOS\*\*\s*\n-+\s*\n', '', cleaned_response, flags=re.MULTILINE)
    
    return cleaned_response.strip()

def convert_markdown_to_text(markdown_report):
    """将Markdown格式报告转换为纯文本格式"""
    import re
    
    # 移除Markdown标记
    text_report = markdown_report
    
    # 移除标题标记
    text_report = re.sub(r'^#+\s*', '', text_report, flags=re.MULTILINE)
    
    # 移除粗体标记
    text_report = re.sub(r'\*\*(.*?)\*\*', r'\1', text_report)
    
    # 移除列表标记，但保留缩进
    text_report = re.sub(r'^\s*[-*]\s*', '  ', text_report, flags=re.MULTILINE)
    
    # 移除代码块标记
    text_report = re.sub(r'```.*?```', '', text_report, flags=re.DOTALL)
    
    # 移除行内代码标记
    text_report = re.sub(r'`(.*?)`', r'\1', text_report)
    
    # 移除链接标记
    text_report = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text_report)
    
    # 清理多余的空行
    text_report = re.sub(r'\n\s*\n\s*\n', '\n\n', text_report)
    
    # 添加分隔线
    text_report = text_report.replace('一、审计概览', '一、审计概览\n' + '='*20)
    text_report = text_report.replace('二、关键漏洞分析', '\n二、关键漏洞分析\n' + '='*20)
    text_report = text_report.replace('三、修复建议', '\n三、修复建议\n' + '='*20)
    text_report = text_report.replace('四、合规性状态', '\n四、合规性状态\n' + '='*20)
    text_report = text_report.replace('五、行动路线图', '\n五、行动路线图\n' + '='*20)
    
    return text_report

def main():
    """主函数"""
    print("正在生成安全审计报告...")
    
    # 解析lynis.log文件
    vulnerabilities = parse_lynis_log("../2025kylin-development/lynis.log")
    
    if not vulnerabilities:
        print("未找到任何漏洞信息，无法生成报告")
        return
    
    # 读取日志文件内容
    try:
        with open("../2025kylin-development/lynis.log", "r", encoding="utf-8") as f:
            log_content = f.read()
    except FileNotFoundError:
        print("错误: 找不到lynis.log文件")
        return
    
    # 生成LLM报告
    markdown_report = generate_llm_report(vulnerabilities, log_content)
    
    # 转换为纯文本格式
    text_report = convert_markdown_to_text(markdown_report)
    
    # 输出Markdown格式报告
    print(markdown_report)
    
    # 输出纯文本格式报告
    print("\n" + "="*50)
    print("安全审计报告（纯文本格式）")
    print("="*50)
    print(text_report)
    
    # 保存Markdown格式报告
    with open("security_audit_report_llm.md", "w", encoding="utf-8") as f:
        f.write(markdown_report)
    
    # 保存纯文本格式报告
    with open("security_audit_report_llm.txt", "w", encoding="utf-8") as f:
        f.write(text_report)
    
    print(f"\n报告已保存到以下文件:")
    print(f"- Markdown格式: security_audit_report_llm.md")
    print(f"- 纯文本格式: security_audit_report_llm.txt")

if __name__ == "__main__":
    main() 